{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4bf69d18",
   "metadata": {},
   "source": [
    "# Vibe-Hacking Tutorial \"My first Hack\" 🤖💻 \n",
    "\n",
    "## Hacking PortSwigger Labs with CAI Framework"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a8ae509-cf6c-4954-a2d9-6f338b881bb5",
   "metadata": {},
   "source": [
    "## 1. Introduction "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ef65e02",
   "metadata": {},
   "source": [
    "We explore how to use the **CAI Python API** to attack PortSwigger Labs. \n",
    "\n",
    "We combine:\n",
    "- CAI’s red-team agent \n",
    "- PortSwigger’s vulnerability lab environments \n",
    "- Large Language Models (LLMs) for autonomous reasoning and payload generation\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6dd86394",
   "metadata": {},
   "source": [
    "## 2. Objectives "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b869988",
   "metadata": {},
   "source": [
    "- Learn how LLMs assist in crafting attacks on web applications\n",
    "- Test vulnerabilities using PortSwigger Web Security Academy\n",
    "- Automate techniques to identify and fix vulnerabilities using CAI\n",
    "- Analyse and defend using LLMs' log interpretation\n",
    "- Learn how to use the CAI Python API"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13bdf02f",
   "metadata": {},
   "source": [
    "### Installing  Python Requirements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "491498de",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "!pip install cai-framework\n",
    "!pip install pandas\n",
    "!pip install selenium\n",
    "!pip install python-dotenv\n",
    "!pip install nest-asyncio"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd0449bf",
   "metadata": {},
   "source": [
    "Because this is a Jupyter Notebook, we need to have Jupyter properly configured and running locally on our computer in order to follow this tutorial successfully.  \n",
    "\n",
    "If you need more information on setting up Jupyter Notebook, please visit [this link](https://docs.jupyter.org/en/latest/install/notebook-classic.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4cf5495-f676-42fd-a784-b475e192cc9f",
   "metadata": {},
   "source": [
    "## 3. Getting the Portswigger Labs Environment \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d62ff80",
   "metadata": {},
   "source": [
    "### 3.1. Importing Required Modules "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3cd0541-05b0-4e95-adbf-963859768528",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Portswigger web scraper bot\n",
    "import utils.portswiggerbot as pb\n",
    "\n",
    "#Other utilities\n",
    "import nest_asyncio\n",
    "import os \n",
    "from dotenv import load_dotenv\n",
    "load_dotenv(override=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c29987ba",
   "metadata": {},
   "source": [
    "### 3.2. Configuration Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5245d3c",
   "metadata": {},
   "source": [
    "### 3.2.1 Setting Up the CAI `.env` File\n",
    "\n",
    "CAI specifies in the documentation that it is necessary to set up an .env file in the same folder as this notebook. For setting up the PortSwigger environment, two additional variables have been added to the template in CAI’s  [.env.example](https://github.com/aliasrobotics/cai/blob/main/.env.example). So, we need to create an .env file with the following variables:\n",
    "\n",
    "```\n",
    "PORTSWIGGER_USERNAME='your-portswigger-email'\n",
    "PORTSWIGGER_PASSWORD='your-portswigger-password'\n",
    "OPENAI_API_KEY='sk-123'\n",
    "ANTHROPIC_API_KEY=''\n",
    "OLLAMA=''\n",
    "PROMPT_TOOLKIT_NO_CPR=1\n",
    "CAI_STREAM=false\n",
    "```\n",
    "✅ Note on PortSwigger: If you don’t have a PortSwigger Labs account, you can create one [here](https://portswigger.net/users/register).\n",
    "\n",
    "✅ Note on CAI: If you need more information on setting up the .env file, check out the [CAI documentation](https://github.com/aliasrobotics/cai/blob/main/README.md#nut_and_bolt-setup-env-file).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b29273ba-10d2-4b1e-9ff2-e1951b45bea3",
   "metadata": {},
   "source": [
    "### 3.2.2 Setting Up the PortSwigger Bot\n",
    "\n",
    "\n",
    "To extract the labs, we start by initialising the PortSwigger bot and loading the username and password from the .env file. Next, we specify the number of labs <b>(N_LABS)</b> to retrieve and the type of vulnerability we want to target <b>(SECTION)</b>. The supported vulnerability sections include:\n",
    "\n",
    "- sql-injection\n",
    "- cross-site-scripting\n",
    "- cross-site-request-forgery-csrf\n",
    "- clickjacking\n",
    "- dom-based-vulnerabilities\n",
    "- cross-origin-resource-sharing-cors\n",
    "- xml-external-entity-xxe-injection\n",
    "- server-side-request-forgery-ssrf\n",
    "- http-request-smuggling\n",
    "- os-command-injection\n",
    "- server-side-template-injection\n",
    "- path-traversal\n",
    "- access-control-vulnerabilities\n",
    "- authentication\n",
    "- websockets\n",
    "- web-cache-poisoning\n",
    "- insecure-deserialization\n",
    "- information-disclosure\n",
    "- business-logic-vulnerabilities\n",
    "- http-host-header-attacks\n",
    "- oauth-authentication\n",
    "- file-upload-vulnerabilities\n",
    "- jwt\n",
    "- essential-skills\n",
    "- prototype-pollution\n",
    "- graphql-api-vulnerabilities\n",
    "- race-conditions\n",
    "- nosql-injection\n",
    "- api-testing\n",
    "- web-llm-attacks\n",
    "- web-cache-deception"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8fa82db-af6b-43e8-bf21-550532748675",
   "metadata": {},
   "outputs": [],
   "source": [
    "SECTION = \"cross-site-scripting\"\n",
    "N_LABS = 3 # If you want to test all the labs in the section, change this to -1\n",
    "USERNAME = os.getenv(\"PORTSWIGGER_USERNAME\")\n",
    "PASSWORD = os.getenv(\"PORTSWIGGER_PASSWORD\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f300d6a7-c689-44c0-85a6-e4bd8f5bea31",
   "metadata": {},
   "source": [
    "### 3.3. Getting  PortSwigger labs\n",
    "We can now extract the lab information by running the following line of code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b90cba8c-ca4d-43fb-a417-862ceb422a0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Initialize bot\n",
    "bot = pb.Bot()\n",
    "\n",
    "#Login to Portswigger Academy\n",
    "bot.login(USERNAME,PASSWORD)\n",
    "\n",
    "#Get labs urls by section or type of lab\n",
    "topics = bot.choose_topic(SECTION)\n",
    "\n",
    "#Get lab metadata\n",
    "labs = [bot.obtain_lab_information(link) for link in topics[0:N_LABS]]\n",
    "\n",
    "#format session cookies\n",
    "cookies = bot.driver.get_cookies()\n",
    "\n",
    "# Define names of essential cookies\n",
    "essential_cookies = {'SessionId','Authenticated_UserVerificationId','t'}\n",
    "# Extract only the essential cookies\n",
    "essential_cookies = [cookie for cookie in cookies if cookie['name'] in essential_cookies]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffb1233d-9952-4437-b469-b48d9fccbdea",
   "metadata": {},
   "source": [
    "## 4. Setting Up the CAI Red Team Agent\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cf5f7ba-8fe2-4864-b3ec-efcbf95c95e0",
   "metadata": {},
   "source": [
    "An Agent is an intelligent component that interacts with an environment. In cybersecurity, an Agent interacts with systems and networks, using peripherals and network interfaces as sensors, reasoning about the gathered information, and then executing network actions as if they were actuators.\n",
    "\n",
    "In CAI, there are different types of predefined agents. You can take a look here:\n",
    "-  [CAI Agents](https://github.com/aliasrobotics/cai/tree/main/src/cai/agents)\n",
    "\n",
    "For this project, we will use the “Red Team Base Agent.” First, we need to specify the LLMs we want to use with it. We are going to use GTP 4.0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "769b427c",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"CAI_MODEL\"] = \"openai/gpt-4o\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50ac55fa-735a-45df-a6c3-71b3acbdfc60",
   "metadata": {},
   "source": [
    "#### 4.1. Prompts explanation\n",
    "\n",
    "CAI Agents work with two types of prompts:\n",
    "\n",
    "##### System Prompt \n",
    "This prompt defines the overall context and instructs the LLMs to assume a specific role. In this case, the model’s goal is to identify and exploit vulnerabilities.\n",
    "\n",
    "✅ Note: When using the built-in Red Team Agent from CAI, the system prompt is already configured for you, so there is no need to set it manually. You can find the default template for the agent here:\n",
    "-  [Red Team Agent System Prompt](https://github.com/aliasrobotics/cai/blob/main/src/cai/prompts/system_red_team_agent.md)\n",
    "\n",
    "##### User Prompt (customizable)\n",
    "\n",
    "This prompt defines the specific task the agent should perform. For PortSwigger Labs, it should include:\n",
    "\n",
    "- The lab title (placeholder: **LAB_TITLE**)\n",
    "- The lab description (placeholder: **LAB_DESCRIPTION**)\n",
    "- The lab environment URL (placeholder: **LAB_ENVIRONMENT_URL**)\n",
    "- Session cookies to authenticate in the lab (placeholder: **COOKIES**)\n",
    "\n",
    "✅ Note: We will use the following user prompt, so it is possible to modify it and interact with or create other types of prompts for LLMs.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e51f40e5-6769-4d2d-9da5-cd8fe9e5f5cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_example = f\"\"\"\n",
    "attack **LAB_ENVIRONMENT_URL**\n",
    "The title is **LAB_TITLE** \n",
    "The description is **LAB_DESCRIPTION** \n",
    "\n",
    "Please consider the following mandatory rules:\n",
    "1. You have to connect to the session with the cookies **COOKIES**\n",
    "2. Do not stop until you try all the commands or exploits you can run and create\n",
    "3. Write the explanation of what you did in an HTML format easy to understand\n",
    "4. Write every payload command and the exploits you create in the format <payloadgen>HERE</payloadgen>\n",
    "5. Do not create files with the payload that you generate because you have to solve the lab by yourself\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c05c5e76",
   "metadata": {},
   "source": [
    "Now that we have explained the agent and prompts, it is time to run our first AI-powered task, but first we must choose the target lab. and format the final user prompt with the data of the lab. In this case, we are going to use the following one:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83b8274c",
   "metadata": {},
   "outputs": [],
   "source": [
    "lab = labs[0]\n",
    "lab"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07d2eaa5",
   "metadata": {},
   "source": [
    "We can take a look at the formatted user prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d9fbf4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "user_prompt = (\n",
    "    prompt_example.replace(\"**LAB_ENVIRONMENT_URL**\", lab['environment_url'])\n",
    "    .replace(\"**LAB_TITLE**\", lab['title'])\n",
    "    .replace(\"**LAB_DESCRIPTION**\", lab['description'])\n",
    "    .replace(\"**COOKIES**\", str(essential_cookies))\n",
    "           )\n",
    "print(user_prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62ee116d",
   "metadata": {},
   "source": [
    "## 5. Attacking PortSwigger Labs with CAI Red Team Agent\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c2ed6b8",
   "metadata": {},
   "source": [
    "It’s time to have some fun hacking PortSwigger Labs! In the code below, we’ve created a helper function that runs the agent asynchronously. Just call the function, and your agent will automatically interact with the target data, searching for vulnerabilities and trying to exploit them"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc92b04a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import after environment setup\n",
    "from cai.sdk.agents import Runner\n",
    "from cai.agents.red_teamer import redteam_agent\n",
    "\n",
    "# Allow async in jupyter notebooks\n",
    "nest_asyncio.apply()\n",
    "\n",
    "\n",
    "#function to run the agent asynchronously\n",
    "async def run_agent(message):\n",
    "    response = await Runner.run(redteam_agent, message)\n",
    "    return response\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ba39128",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = await run_agent(message=user_prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1573d57-eb09-418b-b5f0-a0162cf873be",
   "metadata": {},
   "source": [
    "## 6. Result Analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1e7610a",
   "metadata": {},
   "source": [
    "Now that our CAI agent has done its work, let’s break down the results to understand what it found, how it exploited the target, and what the structure of the final response is. In the following code, we can see the response type and how the output is displayed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75bcbe7d-4d48-4857-91bd-8bfaa3e9147f",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"The tipe of the response is a: {type(response)}\")\n",
    "response"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aafbdb70",
   "metadata": {},
   "source": [
    "### 6.1 Understanding the `RunResult`\n",
    "\n",
    "When we run the CAI Red Team Agent, we get back a RunResult object.  \n",
    "Think of it as a detailed report of  everything the agent did  during its hacking session.  \n",
    "Here’s how to read it:\n",
    "\n",
    "\n",
    "#### `input`\n",
    "\n",
    "This is the user prompt we gave to the agent.  \n",
    "\n",
    "\n",
    "####  `new_items`\n",
    "This shows what the agent  produced during the run. This considers:\n",
    "\n",
    "##### `MessageOutputItem`\n",
    "Contains the following information:\n",
    "  - The agent info  (Red Team Agent)\n",
    "  - The  tools it used (like generic_linux_command to run shell commands, or execute_code to write and run exploits)\n",
    "  - The output, which is a clear HTML report  explaining the methodology used to solve the lab.\n",
    "\n",
    "\n",
    "#### `raw_responses`\n",
    "\n",
    "This is the  raw output  from the agent, showing exactly what text the LLM produced.  \n",
    "\n",
    "\n",
    "#### `final_output`\n",
    "\n",
    "This is the  final, cleaned-up version of the report. \n",
    "\n",
    "In our example, it’s a complete HTML file  that explains:\n",
    "\n",
    "- What the attack did\n",
    "- The payload used:\n",
    "  ```html\n",
    "  <payloadgen><script>alert('XSS Lab')</script></payloadgen>\n",
    "  ```\n",
    "- How to inject it into the lab’s search parameter\n",
    "- The result and security impact"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f105aaf",
   "metadata": {},
   "source": [
    "## 7. Conclusion"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae626b54",
   "metadata": {},
   "source": [
    "In this notebook, we learned how to run a complete attack on the PortSwigger vulnerable lab using the CAI Python API.  \n",
    "\n",
    "Now you are able to:\n",
    "- Launch a CAI **Red Team Agent** \n",
    "- Provide clear instructions, cookies, and payload rules\n",
    "- Let the agent autonomously craft and test exploits\n",
    "- Review the results step by step in a structured `RunResult`\n",
    "\n",
    "This shows the power of combining **LLM-driven reasoning** with real hacking tools.  \n",
    "You can now adapt this approach to test other labs, real applications, or integrate it into your own red teaming workflows.\n",
    "\n",
    "Next, try:\n",
    "- Experimenting with different lab challenges on **PortSwigger Web Security Academy**\n",
    "- Tuning your agent prompts for more advanced payloads\n",
    "- Analysing the results to write better detection and defence rules\n",
    "Happy hacking!\n",
    "\n",
    "# Happy hacking! 🕵️‍♂️💻🔍  "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "test",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
